
import json
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap

# 1. 加载JSON数据
with open('data/math_concept_relations.json', 'r', encoding='utf-8') as f:
    data = json.load(f)

# 2. 创建有向图
G = nx.DiGraph()

# 3. 添加节点并设置属性
level_colors = {'中学': 'lightgreen', '大学': 'lightblue', '人工智能': 'salmon'}
category_shapes = {
    '代数': 's',  # 方形
    '几何': 'd',  # 菱形
    '概率统计': 'o',  # 圆形
    '三角函数': '^',  # 三角形
    '微积分': 'p',  # 五边形
    '线性代数': 'h',  # 六边形
    '概率论与数理统计': '8',  # 八角形
    '离散数学': '*',  # 星形
    '最优化理论': 'X',  # X形
    '矩阵分析': '+',  # 加号
    '信息论': 'D',  # 钻石形
    '数值计算': 'P'  # 加粗的加号
}

for concept in data['concepts']:
    G.add_node(concept['concept_id'], 
               name=concept['concept_name'],
               level=concept['level'],
               category=concept['category'])

# 4. 添加边
for relation in data['relations']:
    G.add_edge(relation['source_concept_id'], 
               relation['target_concept_id'],
               relation_type=relation['relation_type'])

# 5. 设置图形布局
pos = nx.spring_layout(G, k=0.5, iterations=50)

# 6. 绘制节点
node_colors = [level_colors[G.nodes[node]['level']] for node in G.nodes()]
node_shapes = [category_shapes[G.nodes[node]['category']] for node in G.nodes()]

# 7. 绘制图形
plt.figure(figsize=(20, 15))

# 绘制节点
for i, node in enumerate(G.nodes()):
    nx.draw_networkx_nodes(G, pos, 
                          nodelist=[node],
                          node_shape=node_shapes[i],
                          node_size=1500,
                          node_color=node_colors[i])

# 绘制边
edge_colors = ['red' if G.edges[edge]['relation_type'] == 'prerequisite' else 'gray' 
               for edge in G.edges()]
nx.draw_networkx_edges(G, pos, width=1.5, edge_color=edge_colors, arrows=True)

# 绘制标签
labels = {node: G.nodes[node]['name'] for node in G.nodes()}
nx.draw_networkx_labels(G, pos, labels, font_size=10)

# 设置全局字体
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'SimSun', 'Arial Unicode MS']  # 多字体备选
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 8. 添加图例
# 教育阶段图例
legend_elements = [
    plt.Line2D([0], [0], marker='o', color='w', label='中学',
               markerfacecolor='lightgreen', markersize=10),
    plt.Line2D([0], [0], marker='o', color='w', label='大学',
               markerfacecolor='lightblue', markersize=10),
    plt.Line2D([0], [0], marker='o', color='w', label='人工智能',
               markerfacecolor='salmon', markersize=10)
]

# 关系类型图例
legend_elements.extend([
    plt.Line2D([0], [0], color='red', lw=2, label='先修关系'),
    plt.Line2D([0], [0], color='gray', lw=2, label='包含关系')
])

plt.legend(handles=legend_elements, loc='upper right')

# 9. 保存和显示图形
plt.title('数学概念知识图谱', fontsize=16)
plt.tight_layout()
plt.savefig('data/math_concept_graph.png', dpi=300, bbox_inches='tight')
plt.show()